import os
from types import NoneType
from typing import Annotated, List, Literal

from langchain_community.chat_models.zhipuai import ChatZhipuAI
from langchain_core.messages import BaseMessage
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import add_messages, StateGraph, START, END
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel
from typing_extensions import TypedDict

os.environ["ZHIPUAI_API_KEY"] = "97738d4998b8732d707daf91a2b1c56d.2y6VKEuOlidwHDpI"

"""定义State开始"""


class State(TypedDict):
    messages: Annotated[List[BaseMessage], add_messages]
    summary: str
    intention: str
    sop: str
    insurance_policy: list
    surrender_num: str

"""定义State结束"""

"""定义Tools开始"""


# 查询保单
@tool
def insurance_order_search(state: State):
    """查询保单"""
    pass


# 确认退保（用户保留）
@tool
def user_retention(state: State):
    """用户挽留"""
    pass


# 保单退保
@tool
def insurance_order_cancel(state: State):
    """保单退保"""
    pass


# 询问其他事项
@tool
def other_matters(state: State):
    """询问其他事项"""
    pass


class AskHuman(BaseModel):
    """Ask the human a question"""
    question: str


tools = [insurance_order_search, user_retention, insurance_order_cancel, other_matters]
tool_node = ToolNode(tools)
"""定义Tools结束"""

"""定义Model开始"""
model = ChatZhipuAI(
    model='glm-4',
    temperature=0.5
)
"""定义Model结束"""

"""定义Nodes开始"""


class intention(TypedDict):
    intention: Literal["退保", "非退保其他保险业务", "结束会话"]


class sop(TypedDict):
    sop: Literal["二次确认", "确认退保", "确认保单", "执行退保"]


intention_recognition_prompt = """你是一个保险客户专员，能够结合用户的历史消息，识别出最新一条用户消息的核心意图。
历史消息总结:
{summary}

消息上下文
{context}

用户最新的消息
{user_message}
"""


def intention_recognition(state: State):
    prompt = format_prompt(state, intention_recognition_prompt)
    response = model.with_structured_output(intention).invoke(prompt)
    return {"intention": response["intention"]}


def intention_condition(state: State):
    if state["intention"] == '退保':
        return "surrender"
    elif state["intention"] == '非退保其他保险业务':
        return "non_surrender"
    elif state["intention"] == NoneType:
        return "salutation"
    else:
        return "end_session"


def format_prompt(state: State, prompt_template: str) -> str:
    context_messages = []
    if len(state["messages"]) > 1:
        context_messages = state["messages"][0:len(state["messages"]) - 1]
    context = ""
    for item in context_messages:
        context = context + item.type + ":" + item.content + "\n"
    prompt = prompt_template.format(
        summary=state.get("summary", "暂无"),
        context=context,
        user_message=state["messages"][len(state["messages"]) - 1].type + ":" + state["messages"][
            len(state["messages"]) - 1].content)
    return prompt


def non_surrender(state: State):
    return {"messages": [("ai", "好的，我帮你转人工客服进行处理")]}


def salutation(state: State):
    return {"messages": [("ai", "你好，我是水滴保险客服小依，有什么可以帮助到您的")]}


def end_session(state: State):
    return {"messages": [("ai", "好的，那我就先完结此次会话了，祝您生活愉快")]}


def secondary_confirmation(state: State):
    return {"messages": [("ai", "你好，你确认要进行退保操作么？")]}


sop_recognition_prompt = """你是一个保险客户专员，目前用户已经进入到了退保的流程里面，你需要按照如下的SOP步骤，结合用户的历史消息，识别出来下一步应该做什么。
SOP步骤：
1、如果没有对用户进行过是否确认退保的二次确认，请返回：二次确认
2、如果已经进行了二次确认，并且用户最新的消息是确认退保，请返回：确认退保
3、如果用户最新的消息是确认退保的保单，请返回：确认保单

历史消息总结:
{summary}

消息上下文
{context}

用户最新的消息
{user_message}
"""


def sop_recognition(state: State):
    prompt = format_prompt(state, sop_recognition_prompt)
    response = model.with_structured_output(sop).invoke(prompt)
    return {"sop": response["sop"]}


def sop_condition(state: State):
    if state["sop"] == '二次确认':
        return "secondary_confirmation"
    elif state["sop"] == "确认退保":
        return "confirm_surrender"
    elif state["sop"] == "确认保单":
        return "confirm_insurance"
    else:
        return "execute_surrender"


def ask_human(state: State):
    messages = state["messages"]
    return messages

"""定义Nodes结束"""

"""定义WorkFlow开始"""
work_flow = StateGraph(State)
work_flow.add_node("intention_recognition", intention_recognition)
work_flow.add_node("non_surrender", non_surrender)
work_flow.add_node("sop_recognition", sop_recognition)
work_flow.add_node("end_session", end_session)
work_flow.add_node("ask_human", ask_human)
work_flow.add_node("salutation", salutation)
work_flow.add_node("secondary_confirmation", secondary_confirmation)

work_flow.add_edge(START, "intention_recognition")
work_flow.add_conditional_edges("intention_recognition",
                                intention_condition,
                                {"surrender": "sop_recognition",
                                 "non_surrender": "non_surrender",
                                 "salutation": "salutation",
                                 "end_session": "end_session"})
work_flow.add_edge("non_surrender", END)
work_flow.add_edge("end_session", END)
work_flow.add_edge("salutation", END)
work_flow.add_conditional_edges("sop_recognition",
                                sop_condition,
                                {"secondary_confirmation": "secondary_confirmation",
                                 "confirm_surrender": "ask_human",
                                 "confirm_insurance": "ask_human",
                                 "execute_surrender": "ask_human"})
work_flow.add_edge("secondary_confirmation", "ask_human")
work_flow.add_edge("ask_human", "sop_recognition")

graph = work_flow.compile(checkpointer=MemorySaver(), interrupt_before=["ask_human"])

"""定义WorkFlow结束"""

config = {"configurable": {"thread_id": "1"}}
for event in graph.stream({"messages": [("user", "我要退保")]}, config):
    for value in event.values():
        print(value)

graph.update_state(config, {"messages": [("user", "确认退保")]}, as_node="ask_human")
for event in graph.stream(None, config):
    for value in event.values():
        print(value)
